#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <unistd.h>
#include <time.h>
#include <linux/videodev2.h>
#include <xf86drm.h>
#include <xf86drmMode.h>
#include <math.h>

#include "rga_util.h"
#include "util.h"

#define RGA_NPU_INPUT_WIDTH MODEL_INPUT_WIDTH
#define RGA_NPU_INPUT_HEIGHT MODEL_INPUT_HEIGHT

// Map V4L2 formats to RGA formats
static int get_rga_format(int v4l2_format)
{
    switch (v4l2_format)
    {
        case V4L2_PIX_FMT_NV12:
            return RK_FORMAT_YCbCr_420_SP;
        case V4L2_PIX_FMT_NV16:
            return RK_FORMAT_YCbCr_422_SP;
        case V4L2_PIX_FMT_RGB24:
            return RK_FORMAT_RGB_888; // **新增RGB格式支持**
        // Add other format mappings if needed
        default:
            fprintf(stderr, "Unsupported V4L2 format for RGA: %.4s\n", (char *)&v4l2_format);
            return -1;
    }
}

int rga_init(rga_context_t *rga_ctx, rknn_engine_context_t *rknn_ctx)
{
    memset(rga_ctx, 0, sizeof(rga_context_t));

    // 1. Let RKNN allocate its own input buffer for optimal performance.
    // This buffer will be the destination for RGA's preprocessing.
    rga_ctx->npu_input_mem = rknn_create_mem(rknn_ctx->ctx, rknn_ctx->input_attrs[0].size_with_stride);
    RETURN_ON_FAIL(!rga_ctx->npu_input_mem, "Failed to create RKNN input tensor memory\n");

    // 2. Get the file descriptor of the RKNN buffer and import it into RGA.
    // This allows RGA to write directly into the NPU's input memory (Zero-Copy).
    int npu_input_fd = rga_ctx->npu_input_mem->fd;
    im_handle_param_t param =
    {
        .width = RGA_NPU_INPUT_WIDTH,
        .height = RGA_NPU_INPUT_HEIGHT,
        .format = RK_FORMAT_RGB_888
    };
    rga_ctx->npu_input_handle = importbuffer_fd(npu_input_fd, &param);
    RETURN_ON_FAIL(rga_ctx->npu_input_handle <= 0, "Failed to import NPU input buffer to RGA\n");

    printf("RGA module initialized with shared RKNN input buffer.\n");
    return 0;
}

void rga_deinit(rga_context_t *rga_ctx, rknn_engine_context_t *rknn_ctx)
{
    // Buffers must be released in the reverse order of creation/import.
    if (rga_ctx->npu_input_handle)
    {
        releasebuffer_handle(rga_ctx->npu_input_handle);
        rga_ctx->npu_input_handle = 0;
    }
    if (rga_ctx->npu_input_mem)
    {
        rknn_destroy_mem(rknn_ctx->ctx, rga_ctx->npu_input_mem);
        rga_ctx->npu_input_mem = NULL;
    }
}
// 在 rga_util.c 中，用这个版本替换 rga_preprocess_frame
int rga_preprocess_frame(rga_context_t *rga_ctx, int src_dma_fd, int src_width, int src_height, int src_format)
{
    IM_STATUS ret;
    rga_buffer_t src_buf = {0};
    rga_buffer_t dst_buf = {0};
    rga_buffer_handle_t src_handle = 0;

    int rga_fmt = get_rga_format(src_format);
    if (rga_fmt < 0)
        return -1;

    im_handle_param_t src_param = { .width = (uint32_t)src_width, .height = (uint32_t)src_height, .format = (uint32_t)rga_fmt };
    src_handle = importbuffer_fd(src_dma_fd, &src_param);
    if (src_handle <= 0)
    {
        fprintf(stderr, "ERROR: rga_preprocess_frame importbuffer_fd failed for src_dma_fd %d\n", src_dma_fd);
        return -1;
    }

    src_buf = wrapbuffer_handle(src_handle, src_width, src_height, rga_fmt);
    dst_buf = wrapbuffer_handle(rga_ctx->npu_input_handle, MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT, RK_FORMAT_RGB_888);

    float scale = fminf((float)MODEL_INPUT_WIDTH / src_width, (float)MODEL_INPUT_HEIGHT / src_height);
    int resize_w = (int)(src_width * scale);
    int resize_h = (int)(src_height * scale);
    int dst_x = (MODEL_INPUT_WIDTH - resize_w) / 2;
    int dst_y = (MODEL_INPUT_HEIGHT - resize_h) / 2;

    // *** 语法修正 ***
    // 1. 先创建 im_rect 临时变量
    im_rect fill_rect = {0, 0, MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT};
    im_rect src_rect = {0, 0, src_width, src_height};
    im_rect dst_rect = {dst_x, dst_y, resize_w, resize_h};

    // 2. 再将简单的变量传递给宏
    imfill(dst_buf, fill_rect, 0x808080FF);
    ret = improcess(src_buf, dst_buf, (rga_buffer_t)
    {
        0
    }, src_rect, dst_rect, (im_rect)
    {
        0
    }, 0);

    releasebuffer_handle(src_handle);

    if (ret != IM_STATUS_SUCCESS)
    {
        fprintf(stderr, "ERROR: RGA improcess failed: %s\n", imStrError(ret));
        return -1;
    }

    //imsync(-1);
    return 0;
}

// 在 rga_util.c 中，用这个修复后的版本完整替换 rga_draw_results 函数
// 在 rga_util.c 中，用这个最终修复版完整替换 rga_draw_results 函数

int rga_draw_results(int display_dma_fd, int frame_width, int frame_height, detect_result_group_t *results)
{
    if (results == NULL || results->count <= 0) return 0;

    rga_buffer_handle_t display_handle = 0;
    im_handle_param_t param = { .width = (uint32_t)frame_width, .height = (uint32_t)frame_height, .format = RK_FORMAT_YCbCr_420_SP };

    display_handle = importbuffer_fd(display_dma_fd, &param);
    if (display_handle <= 0)
    {
        fprintf(stderr, "ERROR: rga_draw_results importbuffer_fd failed for display_dma_fd %d\n", display_dma_fd);
        return -1;
    }

    rga_buffer_t display_buf = wrapbuffer_handle(display_handle, frame_width, frame_height, RK_FORMAT_YCbCr_420_SP);

    uint32_t yuv_colors[] = {
        (76 << 16) | (84 << 8) | 255,  // Red
        (149 << 16) | (43 << 8) | 21,   // Green
        (29 << 16) | (255 << 8) | 107   // Blue
    };
    
    // *** 关键修复 1: 将线条粗细改为偶数 ***
    const int thickness = 6; 

    for (int i = 0; i < results->count; i++)
    {
        detect_result_t *res = &results->results[i];

        // *** 关键修复 2: 对坐标和尺寸进行2像素对齐 (强制为偶数) ***
        // 使用位运算 `& ~1` 可以高效地将一个数变为小于或等于它自身的偶数
        int aligned_left   = res->left & ~1;
        int aligned_top    = res->top & ~1;
        int aligned_right  = res->right & ~1;
        int aligned_bottom = res->bottom & ~1;

        int box_width  = aligned_right - aligned_left;
        int box_height = aligned_bottom - aligned_top;

        // 确保对齐后的尺寸仍然有效
        if (box_width < thickness * 2 || box_height < thickness * 2)
            continue;
        
        uint32_t color = yuv_colors[res->class_id % 3];

        // 使用对齐后的值进行绘制
        im_rect top_border = {aligned_left, aligned_top, box_width, thickness};
        imfill(display_buf, top_border, color);

        im_rect bottom_border = {aligned_left, aligned_top + box_height - thickness, box_width, thickness};
        imfill(display_buf, bottom_border, color);

        im_rect left_border = {aligned_left, aligned_top + thickness, thickness, box_height - thickness * 2};
        imfill(display_buf, left_border, color);

        im_rect right_border = {aligned_left + box_width - thickness, aligned_top + thickness, thickness, box_height - thickness * 2};
        imfill(display_buf, right_border, color);
    }

    releasebuffer_handle(display_handle);
    //imsync(-1);
    return 0;
}
